package com.xuecheng.aicodenew.ai.rag;


import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

/**
 * 加载rag---标准版
 *
 */
@Configuration
public class RagConfig {

    //注入embedding模型
    @Resource
    private EmbeddingModel qwenEnbeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    //自定义一个检索器
    @Bean
    public ContentRetriever contentRetriever(){
        //rag
        //1.加载原文档
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs");
        //2.文档切割：将原文档按照段落的切割方法进行分割，最大字符1000，重叠最多200个字符
        DocumentByParagraphSplitter paragraphSplitter = new DocumentByParagraphSplitter(1000, 200);
        //3.自定义文档加载器
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(paragraphSplitter)
                //为了提高检索的效率，将文档名添加到每一个分割后的片段（TextSegment）
                .textSegmentTransformer(textSegment -> TextSegment.from(textSegment.metadata().getString("file_name") + "\n" + textSegment.text()
                        , textSegment.metadata()))
                //使用指定的向量模型
                .embeddingModel(qwenEnbeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        //4.加载文档
        ingestor.ingest(documents);

        //5自定义内容查询器
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(qwenEnbeddingModel)
                .embeddingStore(embeddingStore)
                .maxResults(5)//最多 5 个检索结果
                .minScore(0.75)//过滤掉分数小于 0.75 的结果
                .build();
        return contentRetriever;
    }

}
